
# !处理路径导入问题（添加绝对路径）！！！
import sys
import os
CODE_INTERNAL_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), '../../')) # 生成Code文件夹内部对应的绝对路径
sys.path.append(CODE_INTERNAL_PATH)

import pickle
import math
import copy
import matplotlib.pyplot as plt

from utils.draw_data import draw_line_chart

DATA_PATH = "./data/collect_data.pkl"

X_RANGE = [2.1, 3.8]
V_X_RANGE = [-0.2, 0.2]

def read_data(fileName):
  with open(fileName, "rb") as file:
    return pickle.load(file)
  
def find_excellecnt_data(collect_datas):
  excellent_datas = []

  for data in collect_datas:
    max_x = max(data["cavM_filter_x"])
    min_v_x = min(data["cavM_filter_x"])

    # 检查横坐标
    if max_x >= X_RANGE[0] and max_x <= X_RANGE[1]:
      # 检查横向速度
      if min_v_x >= V_X_RANGE[0] and min_v_x <= V_X_RANGE[1]:
        excellent_datas.append(data)

  # 排序(基于横线位置降序)
  sorted(excellent_datas, key=lambda x: max(x["cavM_filter_x"]), reverse=True)

  return excellent_datas

def draw_data(excellent_data):
  hdvH_y = excellent_data['hdvH_y']
  cavM_y = excellent_data['cavM_y']
  cavM_x = excellent_data['cavM_x']
  hdvH_v = excellent_data['hdvH_v']
  cavM_v_y = excellent_data['cavM_v_y']
  cavM_v_x = excellent_data['cavM_v_x']
  cavM_psi = excellent_data['cavM_psi']
  cavM_acc = excellent_data['cavM_acc']
  cavM_delta = excellent_data['cavM_delta']
  collect_reward = excellent_data['collect_reward']
  cav1_y = excellent_data['cav1_y']
  cav1_v = excellent_data['cav1_v']
  hdvF_y = excellent_data['hdvF_y']
  hdvF_v = excellent_data['hdvF_v']
  hdvH_acc = excellent_data['hdvH_acc']
  hdvF_acc = excellent_data['hdvF_acc']

  # 横轴
  data_x = [i for i in range(len(hdvH_y))]

  # 绘制纵向位置
  datas_y = [{"data": hdvH_y, "des": "HDV Position"}, {"data": cavM_y, "des": "M Position"}, {"data": cav1_y, "des": "CAV1 Position"}, {"data": hdvF_y, "des": "HDV-F Position"}]
  labels = {"x": "Time（0.1s）", "y": "纵向位置（m）"}
  title = "纵向位置"
  draw_line_chart(data_x, datas_y, labels, title)

  # 绘制横向位置
  datas_x = [{"data": cavM_x, "des": "M Position"}]
  labels = {"x": "Time（0.1s）", "y": "横向位置（m）"}
  title = "横向位置"
  draw_line_chart(data_x, datas_x, labels, title)

  # 绘制纵向速度
  datas_v_y = [{"data": hdvH_v, "des": "HDV Velocity"}, {"data": cavM_v_y, "des": "M Velocity Y"}, {"data": cav1_v, "des": "CAV1 Velocity"}, {"data": hdvF_v, "des": "HDV-F Velocity"}]
  labels = {"x": "Time（0.1s）", "y": "纵向速度（m/s）"}
  title = "纵向速度"
  draw_line_chart(data_x, datas_v_y, labels, title)

  # 绘制横向速度
  datas_v_x = [{"data": cavM_v_x, "des": "M Velocity X"}]
  labels = {"x": "Time（0.1s）", "y": "横向速度（m/s）"}
  title = "横向速度"
  draw_line_chart(data_x, datas_v_x, labels, title)

  # 绘制航向角
  datas_psi = [{"data": cavM_psi, "des": "M Psi"}]
  labels = {"x": "Time（0.1s）", "y": "航向角（rad）"}
  title = "航向角"
  draw_line_chart(data_x, datas_psi, labels, title)

  # 绘制加速度（缺少一个cav1的加速度）
  datas_acc = [{"data": cavM_acc, "des": "M Acceleration"}, {"data": hdvH_acc, "des": "HDV Acceleration"}, {"data": hdvF_acc, "des": "HDV-F Acceleration"}]
  labels = {"x": "Time（0.1s）", "y": "加速度（m/s²）"}
  title = "加速度"
  draw_line_chart(data_x, datas_acc, labels, title)

  # 绘制转向角
  datas_delta = [{"data": cavM_delta, "des": "M Steering Angle"}]
  labels = {"x": "Time（0.1s）", "y": "转向角（rad）"}
  title = "转向角"
  draw_line_chart(data_x, datas_delta, labels, title)

  # 绘制奖励
  datas_reward = [{"data": collect_reward, "des": "Reward"}]
  labels = {"x": "Time（0.1s）", "y": "奖励（Reward）"}
  title = "奖励"
  draw_line_chart(data_x, datas_reward, labels, title)

def plot_three_lines(x, real, mea, filt, title, ylabel):
  plt.figure(figsize=(8,4))
  plt.plot(x, real, 'g-', label='真实值', linewidth=2)
  if mea is not None:
    plt.plot(x, mea, 'ro', markersize=2, alpha=0.5, label='测量值')
  if filt is not None:
    plt.plot(x, filt, 'b-', label='估计值', linewidth=1)
  plt.xlabel('时间 (s)')
  plt.ylabel(ylabel)
  plt.title(title)
  plt.legend()
  plt.grid(True)
  plt.show()

def draw_kalman_filter_data(excellent_data):
  # 读取真实数据
  hdvH_y = excellent_data['hdvH_y']
  hdvH_v = excellent_data['hdvH_v']
  cavM_y = excellent_data['cavM_y']
  cavM_x = excellent_data['cavM_x']
  cavM_v_y = excellent_data['cavM_v_y']
  cavM_v_x = excellent_data['cavM_v_x']
  cavM_psi = excellent_data['cavM_psi']
  cav1_y = excellent_data['cav1_y']
  cav1_v = excellent_data['cav1_v']
  hdvF_y = excellent_data['hdvF_y']
  hdvF_v = excellent_data['hdvF_v']

  # 读取测量数据
  hdvH_mea_y =  excellent_data["hdvH_mea_y"]
  hdvH_mea_v = excellent_data["hdvH_mea_v"]
  cavM_mea_y = excellent_data["cavM_mea_y"]
  cavM_mea_x = excellent_data["cavM_mea_x"]
  cavM_mea_psi = excellent_data["cavM_mea_psi"]
  cavM_mea_v = excellent_data["cavM_mea_v"]
  cavM_mea_v_y = excellent_data["cavM_mea_v_y"]
  cavM_mea_v_x = excellent_data["cavM_mea_v_x"]
  cav1_mea_y = excellent_data["cav1_mea_y"]
  cav1_mea_v = excellent_data["cav1_mea_v"]
  hdvF_mea_y = excellent_data["hdvF_mea_y"]
  hdvF_mea_v = excellent_data["hdvF_mea_v"]

  # 读取滤波数据
  hdvH_filter_y = excellent_data["hdvH_filter_y"]
  hdvH_filter_v = excellent_data["hdvH_filter_v"]
  hdvH_filter_covariances_y = excellent_data["hdvH_filter_covariances_y"]
  hdvH_filter_covariances_v = excellent_data["hdvH_filter_covariances_v"]
  cavM_filter_y = excellent_data["cavM_filter_y"]
  cavM_filter_x = excellent_data["cavM_filter_x"]
  cavM_filter_psi = excellent_data["cavM_filter_psi"]
  cavM_filter_v = excellent_data["cavM_filter_v"]
  cavM_filter_v_y = excellent_data["cavM_filter_v_y"]
  cavM_filter_v_x = excellent_data["cavM_filter_v_x"]
  cavM_filter_covariances_y = excellent_data["cavM_filter_covariances_y"]
  cavM_filter_covariances_x = excellent_data["cavM_filter_covariances_x"]
  cavM_filter_covariances_psi = excellent_data["cavM_filter_covariances_psi"]
  cavM_filter_covariances_v = excellent_data["cavM_filter_covariances_v"]
  cav1_filter_y = excellent_data["cav1_filter_y"]
  cav1_filter_v = excellent_data["cav1_filter_v"]
  cav1_filter_covariances_y = excellent_data["cav1_filter_covariances_y"]
  cav1_filter_covariances_v = excellent_data["cav1_filter_covariances_v"]
  hdvF_filter_y = excellent_data["hdvF_filter_y"]
  hdvF_filter_v = excellent_data["hdvF_filter_v"]
  hdvF_filter_covariances_y = excellent_data["hdvF_filter_covariances_y"]
  hdvF_filter_covariances_v = excellent_data["hdvF_filter_covariances_v"]

  # 绘制
  data_x = [i for i in range(len(hdvH_y))]

  # hdvH的y位置
  plot_three_lines(data_x, hdvH_y, hdvH_mea_y, hdvH_filter_y, 'hdvH的y位置', 'Y位置 (m)')
  print("hdvH_y", hdvH_y)
  print("hdvH_filter_y", hdvH_filter_y)

  # cavM的y位置
  plot_three_lines(data_x, cavM_y, cavM_mea_y, cavM_filter_y, 'cavM的y位置', 'Y位置 (m)')

  # cav1的y位置
  plot_three_lines(data_x, cav1_y, cav1_mea_y, cav1_filter_y, 'cav1的y位置', 'Y位置 (m)')

  # hdvF的y位置
  plot_three_lines(data_x, hdvF_y, hdvF_mea_y, hdvF_filter_y, 'hdvF的y位置', 'Y位置 (m)')

  # cavM的x位置
  plot_three_lines(data_x, cavM_x, cavM_mea_x, cavM_filter_x, 'cavM的x位置', 'Y位置 (m)')

  # hdvH的v
  plot_three_lines(data_x, hdvH_v, hdvH_mea_v, hdvH_filter_v, 'hdvH的v', 'Y位置 (m)')

  # cavM的v_y
  plot_three_lines(data_x, cavM_v_y, cavM_mea_v_y, cavM_filter_v_y, 'cavM的v_y', 'Y位置 (m)')
  
  # cavM的v_x
  plot_three_lines(data_x, cavM_v_x, cavM_mea_v_x, cavM_filter_v_x, 'cavM的v_x', 'Y位置 (m)')

  # cav1的v
  plot_three_lines(data_x, cav1_v, cav1_mea_v, cav1_filter_v, 'cav1的v', 'Y位置 (m)')

  # hdvF的v
  plot_three_lines(data_x, hdvF_v, hdvF_mea_v, hdvF_filter_v, 'hdvF的v', 'Y位置 (m)')

  # cavM的psi
  plot_three_lines(data_x, cavM_psi, cavM_mea_psi, cavM_filter_psi, 'cavM的psi', 'Y位置 (m)')

  # hdvH的协方差
  datas_hdvH_covariances = [{"data": hdvH_filter_covariances_y, "des": "hdvH的y协方差"}, {"data": hdvH_filter_covariances_v, "des": "hdvH的v协方差"} ]
  labels = {"x": "Time（0.1s）", "y": "协方差"}
  title = "hdvH的协方差"
  draw_line_chart(data_x, datas_hdvH_covariances, labels, title)

  # cavM的协方差
  datas_cavM_covariances = [{"data": cavM_filter_covariances_y, "des": "cavM的y协方差"}, {"data": cavM_filter_covariances_x, "des": "cavM的x协方差"}, {"data": cavM_filter_covariances_psi, "des": "cavM的psi协方差"}, ]
  labels = {"x": "Time（0.1s）", "y": "协方差"}
  title = "cavM的协方差"
  draw_line_chart(data_x, datas_cavM_covariances, labels, title)

  # cav1的协方差
  datas_cav1_covariances = [{"data": cav1_filter_covariances_y, "des": "cav1的y协方差"}, {"data": cav1_filter_covariances_v, "des": "cav1的v协方差"} ]
  labels = {"x": "Time（0.1s）", "y": "协方差"}
  title = "cav1的协方差"
  draw_line_chart(data_x, datas_cav1_covariances, labels, title)
  
  # hdvF的协方差
  datas_hdvF_covariances = [{"data": hdvF_filter_covariances_y, "des": "hdvF的y协方差"}, {"data": hdvF_filter_covariances_v, "des": "hdvF的v协方差"} ]
  labels = {"x": "Time（0.1s）", "y": "协方差"}
  title = "hdvF的协方差"
  draw_line_chart(data_x, datas_hdvF_covariances, labels, title)

if __name__ == "__main__":
  # 加载数据
  collect_datas = read_data(DATA_PATH)
  print("Data: ", len(collect_datas))

  # 挑选数据
  excellent_datas = find_excellecnt_data(collect_datas)
  if (len(excellent_datas) == 0):
    print("没有符合条件的数据！！！")
    sys.exit()
  print("优秀数据: ", len(excellent_datas))

  # 绘制数据
  draw_idx = 0
  draw_data(excellent_datas[draw_idx])
  # draw_kalman_filter_data(excellent_datas[draw_idx])